import numpy as np
import cv2
import os
import sys
from PIL import Image
from tqdm import tqdm
import imageio
from imageio import imread,imsave
import skimage

'''
cv2读取视频失败，使用imageio读取正常
'''
def extractFrame(video_file, saveDir, rate):
    videoname = video_file.split("\\")[-1].split(".")[0]
    # 可以选择解码工具
    vid = imageio.get_reader(video_file, 'ffmpeg')
    i = 0
    for index, img in enumerate(vid):
        i += 1
        # rate = 16  # 16帧抽一张
        if (i % rate != 0):
            continue
        fileName1 = videoname + "_" + str(i) + ".jpg"
        data = video_file.split("\\")[-2]
        fileName = data + "_" + fileName1
        imsave(saveDir+ fileName, img)


def cut(video_file, target_dir, warnflag, rate):
    videoname = video_file.split("\\")[-1].split(".")[0]
    cap = cv2.VideoCapture(video_file)  # 获取到一个视频
    isOpened = cap.isOpened  # 判断是否打开

    # 视频属性
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取原视频的宽
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取原视频的搞
    fps = int(cap.get(cv2.CAP_PROP_FPS))  # 帧率
    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))  # 视频的编码


    # 为单张视频，以视频名称所谓文件名，创建文件夹
    # temp = os.path.split(video_file)[-1]
    # dir_name = temp.split('.')[0]
    dir_name = ""

    single_pic_store_dir = os.path.join(target_dir, dir_name)
    if not os.path.exists(single_pic_store_dir):
        os.makedirs(single_pic_store_dir)


    i = 0
    while isOpened:
        i += 1
        # print(i)
        (flag, frame) = cap.read()  # 读取一张图像

        # 隔帧抽取
        # rate = 16  # 16帧抽一张
        if (i % rate != 0):
            continue
        fileName1 = videoname + "_" + str(i) + "_" + str(warnflag) + ".jpg"
        # data = video_file.split("\\")[-3].split("-")[1]
        # data = video_file.split("\\")[-2]
        # fileName = data + "_" + fileName1
        # fileName = str(i) + ".jpg"
        if (flag == True):
            # 以下三行 进行 旋转
            #frame = np.rot90(frame, -1)
            #print(fileName)
            # 设置保存路径
            save_path = os.path.join(single_pic_store_dir, fileName1)
            #print(save_path)
            # cv2.imshow("", frame)
            # cv2.waitKey(0)
            # cv2.imwrite(save_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 100])
            cv2.imencode('.jpg', frame)[1].tofile(save_path)  # 存储失败
            # image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            # image.show()
            # cv2.waitKey(0)
            # image.save(save_path)
            #print(res)
        else:
            break

    return single_pic_store_dir

def fukang(videoDir, saveDir):
    i = 0
    do_list = ["抽烟报警(副驾)", "抽烟报警(后排)", "抽烟报警(主驾)"]
    # do_list = ["接打电话报警(主驾)"]
    for root, dirs, files in os.walk(videoDir):
        for dir in dirs:
            videoDir1 = os.path.join(root, dir)
            for root1, dirs1, files1 in os.walk(videoDir1):
                for dir1 in dirs1:
                    if dir1 in do_list:
                        # 标记报警类型
                        '''
                        43:'抽烟报警(后排)',  53:'抽烟报警(副驾)', 63:'抽烟报警(主驾)',
                            41:'疲劳报警(后排)', 51:'疲劳报警(副驾)', 61:'疲劳驾驶报警(主驾)',
                            62:'接打电话报警(主驾)'
                        '''
                        if dir1 == "抽烟报警(副驾)":
                            warnflag = 53
                        elif dir1 == "抽烟报警(后排)":
                            warnflag = 43
                        elif dir1 == "抽烟报警(主驾)":
                            warnflag = 63
                        elif dir1 == "接打电话报警(主驾)":
                            warnflag = 62
                        videoDir2 = os.path.join(root1, dir1)
                        print(videoDir2)
                        for file in tqdm(os.listdir(videoDir2)):
                            i += 1
                            videofile = os.path.join(videoDir2, file)
                            # print(videofile)
                            LPR = file.split("_")[0]
                            # savePath = videofile.split(".")[0].replace("\\video\\", "\\frame-16\\") + "\\"
                            savePath = saveDir + dir1 + "\\" + LPR + "\\"
                            # print(savePath)
                            cut(videofile, savePath, warnflag)
    print(i)


def zhongshiyou(videoDir, saveDir):
    i = 0
    dir_list = ["20220704"]
    for root, dirs, files in os.walk(videoDir):
        for dir in dirs:
            if dir in dir_list:
                videoPath = os.path.join(videoDir, dir)
                for file in tqdm(os.listdir(videoPath)):
                    i += 1
                    if i < 7:
                        continue
                    videofile = os.path.join(videoPath, file)
                    # print(savePath)
                    # cut(videofile, saveDir, 00, 16)
                    extractFrame(videofile, saveDir, 64)

if __name__ == '__main__':
    videoDir = "Z:\\电科未检测到驾驶员误报数据_20240108"
    saveDir = "F:\\电科未检测到驾驶员视频附件_拆帧\\"
    for v in tqdm(os.listdir(videoDir)):
        if v.endswith("mp4"):
            cut(os.path.join(videoDir, v), saveDir, 00, 5)


